from flask import Flask, request, jsonify
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad, unpad
import sqlite3
import base64
import os

app = Flask(__name__)

# AES 密钥和初始化向量（保证安全）
key = b'1234567890123456'
iv = b'1234567890123456'

db_path = 'notes.db'

# 创建数据库
def create_db():
    if not os.path.exists(db_path):
        conn = sqlite3.connect(db_path)
        c = conn.cursor()
        c.execute('''CREATE TABLE notes (id INTEGER PRIMARY KEY, name TEXT, content BLOB)''')
        conn.commit()
        conn.close()

# 加密数据
def encrypt_data(data):
    cipher = AES.new(key, AES.MODE_CBC, iv)
    encrypted_data = cipher.encrypt(pad(data.encode(), AES.block_size))
    return base64.b64encode(encrypted_data).decode()

# 解密数据
def decrypt_data(encrypted_data):
    cipher = AES.new(key, AES.MODE_CBC, iv)
    encrypted_data = base64.b64decode(encrypted_data)
    decrypted_data = unpad(cipher.decrypt(encrypted_data), AES.block_size).decode()
    return decrypted_data

# 创建数据库
create_db()

@app.route('/create_note', methods=['POST'])
def create_note():
    data = request.get_json()
    name = data.get('name')
    password = data.get('password')
    content = data.get('content')

    if not name or not password or not content:
        return jsonify({'error': 'Missing data'}), 400

    encrypted_content = encrypt_data(content)

    conn = sqlite3.connect(db_path)
    c = conn.cursor()
    c.execute('INSERT INTO notes (name, content) VALUES (?, ?)', (name, encrypted_content))
    conn.commit()
    conn.close()

    return jsonify({'message': 'Note created successfully'}), 201

@app.route('/view_note', methods=['GET'])
def view_note():
    name = request.args.get('name')

    if not name:
        return jsonify({'error': 'Note name required'}), 400

    conn = sqlite3.connect(db_path)
    c = conn.cursor()
    c.execute('SELECT content FROM notes WHERE name = ?', (name,))
    row = c.fetchone()
    conn.close()

    if row:
        decrypted_content = decrypt_data(row[0])
        return jsonify({'content': decrypted_content}), 200
    else:
        return jsonify({'error': 'Note not found'}), 404

@app.route('/delete_note', methods=['DELETE'])
def delete_note():
    name = request.get_json().get('name')

    if not name:
        return jsonify({'error': 'Note name required'}), 400

    conn = sqlite3.connect(db_path)
    c = conn.cursor()
    c.execute('DELETE FROM notes WHERE name = ?', (name,))
    conn.commit()
    conn.close()

    return jsonify({'message': 'Note deleted successfully'}), 200

if __name__ == '__main__':
    app.run(port=5055)
